import numpy as np
import pandas as pd
import scipy.io as sio
import scipy.signal as signal
import pywt
import os
from tqdm import tqdm
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.feature_selection import SelectKBest, f_classif
import warnings
warnings.filterwarnings('ignore')

class DataPreprocessor:
    def __init__(self, data_path, target_sampling_rate=12000):
        self.data_path = data_path
        self.target_fs = target_sampling_rate
        self.scaler = StandardScaler()
        
    def load_mat_files(self, file_list):
        """批量加载MAT文件数据"""
        all_data = []
        file_info = []
        
        for file_name in tqdm(file_list, desc="Loading MAT files"):
            try:
                file_path = os.path.join(self.data_path, file_name)
                data = sio.loadmat(file_path)
                
                # 提取信号数据
                de_signal = self._extract_signal(data, 'DE')
                fe_signal = self._extract_signal(data, 'FE')
                ba_signal = self._extract_signal(data, 'BA')
                rpm = self._extract_rpm(data)
                
                # 提取文件信息
                file_type = self._parse_file_type(file_name)
                fault_type = self._parse_fault_type(file_name)
                fault_size = self._parse_fault_size(file_name)
                load_condition = self._parse_load_condition(file_name)
                
                all_data.append({
                    'file_name': file_name,
                    'de_signal': de_signal,
                    'fe_signal': fe_signal,
                    'ba_signal': ba_signal,
                    'rpm': rpm,
                    'file_type': file_type,
                    'fault_type': fault_type,
                    'fault_size': fault_size,
                    'load_condition': load_condition,
                    'sampling_rate': self._get_sampling_rate(file_name)
                })
                
                file_info.append({
                    'file_name': file_name,
                    'fault_type': fault_type,
                    'fault_size': fault_size,
                    'load_condition': load_condition,
                    'rpm': rpm,
                    'signal_length': len(de_signal)
                })
                
            except Exception as e:
                print(f"Error loading {file_name}: {str(e)}")
                continue
        
        return all_data, pd.DataFrame(file_info)
    
    def _extract_signal(self, data, signal_type):
        """提取特定类型的信号"""
        for key in data.keys():
            if signal_type in key and 'time' in key:
                return data[key].flatten()
        return None
    
    def _extract_rpm(self, data):
        """提取RPM信息"""
        for key in data.keys():
            if 'RPM' in key:
                return data[key][0][0]
        return None
    
    def _parse_file_type(self, file_name):
        """解析文件类型"""
        if 'N' in file_name.split('_')[0]:
            return 'Normal'
        else:
            return 'Fault'
    
    def _parse_fault_type(self, file_name):
        """解析故障类型"""
        prefix = file_name.split('_')[0]
        if 'OR' in prefix:
            return 'Outer Race'
        elif 'IR' in prefix:
            return 'Inner Race'
        elif 'B' in prefix:
            return 'Ball'
        elif 'N' in prefix:
            return 'Normal'
        else:
            return 'Unknown'
    
    def _parse_fault_size(self, file_name):
        """解析故障尺寸"""
        parts = file_name.split('_')
        if len(parts) > 1:
            size_str = parts[0][1:]  # 去掉第一个字母
            if size_str.isdigit():
                return int(size_str) / 1000  # 转换为英寸
        return 0.0
    
    def _parse_load_condition(self, file_name):
        """解析载荷条件"""
        parts = file_name.split('_')
        if len(parts) > 1:
            return int(parts[1])
        return 0
    
    def _get_sampling_rate(self, file_name):
        """获取采样率"""
        if '48k' in file_name.lower():
            return 48000
        elif '12k' in file_name.lower():
            return 12000
        else:
            return 12000  # 默认值

# 使用示例
data_path = "path/to/your/data"
preprocessor = DataPreprocessor(data_path)
file_list = [f for f in os.listdir(data_path) if f.endswith('.mat')]
all_data, file_info = preprocessor.load_mat_files(file_list[:50])  # 加载前50个文件作为示例
